module axi_stream_insert_header #(
    parameter DATA_WD = 32,
    parameter DATA_BYTE_WD = DATA_WD / 8,
    parameter BYTE_CNT_WD = $clog2(DATA_BYTE_WD)
) (
    input    clk,
    input    rst_n,

    // 原始数据输入
    input                   valid_in,
    input [DATA_WD-1:0]     data_in,
    input [DATA_BYTE_WD-1:0] keep_in,
    input                   last_in,
    output                  ready_in,

    // 输出接口
    output                  valid_out,
    output [DATA_WD-1:0]    data_out,
    output [DATA_BYTE_WD-1:0] keep_out,
    output                  last_out,
    input                   ready_out,

    // Header插入接口
    input                   valid_insert,
    input [DATA_WD-1:0]     data_insert,
    input [DATA_BYTE_WD-1:0] keep_insert,
    input [BYTE_CNT_WD-1:0] byte_insert_cnt,
    output                  ready_insert
);

//----------------------------------------------------------
// 核心寄存器定义
//----------------------------------------------------------
 reg [DATA_WD-1:0]      header_buf = 0;
 reg [DATA_BYTE_WD-1:0] header_keep = 0;
 reg [BYTE_CNT_WD:0]    header_remain = 0;
 reg                    header_active = 0;
 reg [DATA_WD-1:0]      data_buf = 0;
 reg [DATA_BYTE_WD-1:0] data_keep = 0;
 reg                    data_last = 0;
 reg                    data_ready = 1;

//----------------------------------------------------------
// Header处理逻辑（对齐校验）
//----------------------------------------------------------
always @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
        header_buf    <= 0;
        header_keep   <= 0;
        header_remain <= 0;
        header_active <= 0;
    end else begin
        if (valid_insert && ready_insert) begin
            if (keep_insert & ((1 << byte_insert_cnt)-1) != keep_insert) begin
                $display("[ERROR] Invalid keep_insert alignment at %t", $time);
            end
            
            header_buf    <= data_insert << (8 * (DATA_BYTE_WD - byte_insert_cnt));
            header_keep   <= keep_insert << (DATA_BYTE_WD - byte_insert_cnt);
            header_remain <= byte_insert_cnt;
            header_active <= 1'b1;
        end else if (header_active) begin
            if (valid_out && ready_out) begin
                header_buf    <= header_buf >> 8;
                header_keep   <= {1'b0, header_keep[DATA_BYTE_WD-1:1]};
                header_remain <= header_remain - 1;
                if (header_remain == 1) header_active <= 0;
            end
        end
    end
end

//----------------------------------------------------------
// 数据处理逻辑
//----------------------------------------------------------
always @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
        data_buf   <= 0;
        data_keep  <= 0;
        data_last  <= 0;
        data_ready <= 1;
    end else begin
        if (valid_in && ready_in) begin
            // 严格协议检查
            if (!last_in && (keep_in != {DATA_BYTE_WD{1'b1}})) begin
                data_keep <= {DATA_BYTE_WD{1'b1}};  // 强制修正
                $display("[WARN]  Non-last beat keep_in violation at %t", $time);
            end else begin
                data_keep <= keep_in;
            end
            data_buf  <= data_in;
            data_last <= last_in;
        end
        
        // 优化反压逻辑
        data_ready <= !header_active && (ready_out || (!valid_out && !data_last));
    end
end

//----------------------------------------------------------
// 输出逻辑（末尾数据保持）
//----------------------------------------------------------
assign valid_out = header_active ? (header_remain > 0) : 
                  (valid_in || (data_last && !ready_out));

assign data_out  = header_active ? 
                 {header_buf[DATA_WD-1:8], 8'h00} : 
                 data_buf;

assign keep_out  = header_active ? 
                 {{DATA_BYTE_WD-1{1'b0}}, header_keep[DATA_BYTE_WD-1]} : 
                 (data_last ? data_keep : {DATA_BYTE_WD{1'b1}});

assign last_out  = !header_active && data_last && ready_out;

//----------------------------------------------------------
// 反压信号生成（防止死锁）
//----------------------------------------------------------
assign ready_insert = !header_active && !(valid_out && !ready_out);
assign ready_in     = data_ready;

endmodule